from enum import Enum, auto
from typing import List, Optional

import os
from typing import TextIO, Union


class StringFileWrapper:
    def __init__(self, fd: TextIO, CHUNK_LENGTH: int) -> None:

        self.fd = fd
        self.length: int = 0
        self.buffers: dict[int, str] = {}
        if not CHUNK_LENGTH or CHUNK_LENGTH < 2:
            CHUNK_LENGTH = 1_000_000
        self.buffer_length = CHUNK_LENGTH

    def get_buffer(self, index: int) -> str:

        if self.buffers.get(index) is None:
            self.fd.seek(index * self.buffer_length)
            self.buffers[index] = self.fd.read(self.buffer_length)
            if len(self.buffers) > max(2, 2_000_000 / self.buffer_length):
                oldest_key = next(iter(self.buffers))
                if oldest_key != index:
                    self.buffers.pop(oldest_key)
        return self.buffers[index]

    def __getitem__(self, index: Union[int, slice]) -> str:

        if isinstance(index, slice):
            buffer_index = index.start // self.buffer_length
            buffer_end = index.stop // self.buffer_length
            if buffer_index == buffer_end:
                return self.get_buffer(buffer_index)[
                    index.start % self.buffer_length : index.stop % self.buffer_length
                ]
            else:
                start_slice = self.get_buffer(buffer_index)[
                    index.start % self.buffer_length :
                ]
                end_slice = self.get_buffer(buffer_end)[
                    : index.stop % self.buffer_length
                ]
                middle_slices = [
                    self.get_buffer(i) for i in range(buffer_index + 1, buffer_end)
                ]
                return start_slice + "".join(middle_slices) + end_slice
        else:
            buffer_index = index // self.buffer_length
            return self.get_buffer(buffer_index)[index % self.buffer_length]

    def __len__(self) -> int:

        if self.length < 1:
            current_position = self.fd.tell()
            self.fd.seek(0, os.SEEK_END)
            self.length = self.fd.tell()
            self.fd.seek(current_position)
        return self.length

    def __setitem__(self, index: Union[int, slice], value: str) -> None:

        if isinstance(index, slice):
            start = index.start or 0
        else:
            start = index or 0

        if start < 0:
            start += len(self)

        current_position = self.fd.tell()
        self.fd.seek(start)
        self.fd.write(value)
        self.fd.seek(current_position)


class ContextValues(Enum):
    OBJECT_KEY = auto()
    OBJECT_VALUE = auto()
    ARRAY = auto()


class JsonContext:
    def __init__(self) -> None:
        self.context: List[ContextValues] = []
        self.current: Optional[ContextValues] = None
        self.empty: bool = True

    def set(self, value: ContextValues) -> None:

        self.context.append(value)
        self.current = value
        self.empty = False

    def reset(self) -> None:

        try:
            self.context.pop()
            self.current = self.context[-1]
        except IndexError:
            self.current = None
            self.empty = True



from typing import Any, Dict, List, Literal, Optional, TextIO, Tuple, Union


JSONReturnType = Union[Dict[str, Any], List[Any], str, float, int, bool, None]


class JSONParser:
    STRING_DELIMITERS = ['"', "'", "“", "”"]

    def __init__(
        self,
        json_str: Union[str, StringFileWrapper],
        json_fd: Optional[TextIO],
        logging: Optional[bool],
        json_fd_chunk_length: int = 0,
    ) -> None:
        self.json_str: Union[str, StringFileWrapper] = json_str
        if json_fd:
            self.json_str = StringFileWrapper(json_fd, json_fd_chunk_length)
        self.index: int = 0
        self.context = JsonContext()

        self.logging = logging
        if logging:
            self.logger: List[Dict[str, str]] = []
            self.log = self._log
        else:
            self.log = lambda *args, **kwargs: None

    def parse(
        self,
    ) -> Union[JSONReturnType, Tuple[JSONReturnType, List[Dict[str, str]]]]:
        json = self.parse_json()
        if self.index < len(self.json_str):
            self.log(
                "The parser returned early, checking if there's more json elements",
            )
            json = [json]
            last_index = self.index
            while self.index < len(self.json_str):
                j = self.parse_json()
                if j != "":
                    json.append(j)
                if self.index == last_index:
                    self.index += 1
                last_index = self.index
            if len(json) == 1:
                self.log(
                    "There were no more elements, returning the element without the array",
                )
                json = json[0]
        if self.logging:
            return json, self.logger
        else:
            return json

    def parse_json(
        self,
    ) -> JSONReturnType:
        while True:
            char = self.get_char_at()
            if char is False:
                return ""
            elif char == "{":
                self.index += 1
                return self.parse_object()
            elif char == "[":
                self.index += 1
                return self.parse_array()
            elif self.context.current == ContextValues.OBJECT_VALUE and char == "}":
                self.log(
                    "At the end of an object we found a key with missing value, skipping",
                )
                return ""
            elif not self.context.empty and (
                char in self.STRING_DELIMITERS or char.isalpha()
            ):
                return self.parse_string()
            elif not self.context.empty and (
                char.isdigit() or char == "-" or char == "."
            ):
                return self.parse_number()
            else:
                self.index += 1

    def parse_object(self) -> Dict[str, JSONReturnType]:
        obj = {}
        while (self.get_char_at() or "}") != "}":

            self.skip_whitespaces_at()

            if (self.get_char_at() or "") == ":":
                self.log(
                    "While parsing an object we found a : before a key, ignoring",
                )
                self.index += 1

            self.context.set(ContextValues.OBJECT_KEY)

            self.skip_whitespaces_at()

            rollback_index = self.index

            key = ""
            while self.get_char_at():
                rollback_index = self.index
                key = str(self.parse_string())

                if key != "" or (key == "" and self.get_char_at() == ":"):
                    break
            if ContextValues.ARRAY in self.context.context and key in obj:
                self.log(
                    "While parsing an object we found a duplicate key, closing the object here and rolling back the index",
                )
                self.index = rollback_index - 1
                self.json_str = (
                    self.json_str[: self.index + 1]
                    + "{"
                    + self.json_str[self.index + 1 :]
                )
                break

            self.skip_whitespaces_at()

            if (self.get_char_at() or "}") == "}":
                continue

            self.skip_whitespaces_at()

            if (self.get_char_at() or "") != ":":
                self.log(
                    "While parsing an object we missed a : after a key",
                )

            self.index += 1
            self.context.reset()
            self.context.set(ContextValues.OBJECT_VALUE)
            value = self.parse_json()

            self.context.reset()
            obj[key] = value

            if (self.get_char_at() or "") in [",", "'", '"']:
                self.index += 1

            self.skip_whitespaces_at()

        self.index += 1
        return obj

    def parse_array(self) -> List[JSONReturnType]:
        arr = []
        self.context.set(ContextValues.ARRAY)
        char = self.get_char_at()
        while char and char not in ["]", "}"]:
            self.skip_whitespaces_at()
            value = self.parse_json()

            if value == "":
                break

            if value == "..." and self.get_char_at(-1) == ".":
                self.log(
                    "While parsing an array, found a stray '...'; ignoring it",
                )
            else:
                arr.append(value)

            char = self.get_char_at()
            while char and (char.isspace() or char == ","):
                self.index += 1
                char = self.get_char_at()

        char = self.get_char_at()
        if char and char != "]":
            self.log(
                "While parsing an array we missed the closing ], adding it back",
            )
            self.index -= 1

        self.index += 1
        self.context.reset()
        return arr

    def parse_string(self) -> Union[str, bool, None]:

        missing_quotes = False
        doubled_quotes = False
        lstring_delimiter = rstring_delimiter = '"'

        char = self.get_char_at()
        while char and char not in self.STRING_DELIMITERS and not char.isalnum():
            self.index += 1
            char = self.get_char_at()

        if not char:
            return ""

        if char == "'":
            lstring_delimiter = rstring_delimiter = "'"
        elif char == "“":
            lstring_delimiter = "“"
            rstring_delimiter = "”"
        elif char.isalnum():
            if (
                char.lower() in ["t", "f", "n"]
                and self.context.current != ContextValues.OBJECT_KEY
            ):
                value = self.parse_boolean_or_null()
                if value != "":
                    return value
            self.log(
                "While parsing a string, we found a literal instead of a quote",
            )
            self.log(
                "While parsing a string, we found no starting quote. Will add the quote back",
            )
            missing_quotes = True

        if not missing_quotes:
            self.index += 1

        self.skip_whitespaces_at()
        if self.get_char_at() in self.STRING_DELIMITERS:
            if self.get_char_at() == lstring_delimiter:
                if (
                    self.context.current == ContextValues.OBJECT_KEY
                    and self.get_char_at(1) == ":"
                ):
                    self.index += 1
                    return ""
                if self.get_char_at(1) == lstring_delimiter:
                    self.log(
                        "While parsing a string, we found a doubled quote and then a quote again, ignoring it",
                    )
                    return ""
                i = self.skip_to_character(character=rstring_delimiter, idx=1)
                next_c = self.get_char_at(i)
                if next_c and (self.get_char_at(i + 1) or "") == rstring_delimiter:
                    self.log(
                        "While parsing a string, we found a valid starting doubled quote",
                    )
                    doubled_quotes = True
                    self.index += 1
                else:
                    i = self.skip_whitespaces_at(idx=1, move_main_index=False)
                    next_c = self.get_char_at(i)
                    if next_c in self.STRING_DELIMITERS + ["{", "["]:
                        self.log(
                            "While parsing a string, we found a doubled quote but also another quote afterwards, ignoring it",
                        )
                        self.index += 1
                        return ""
                    elif next_c not in [",", "]", "}"]:
                        self.log(
                            "While parsing a string, we found a doubled quote but it was a mistake, removing one quote",
                        )
                        self.index += 1
            else:
                i = self.skip_to_character(character=rstring_delimiter, idx=1)
                next_c = self.get_char_at(i)
                if not next_c:
                    self.log(
                        "While parsing a string, we found a quote but it was a mistake, ignoring it",
                    )
                    return ""

        string_acc = ""

        char = self.get_char_at()
        unmatched_delimiter = False
        while char and char != rstring_delimiter:
            if (
                missing_quotes
                and self.context.current == ContextValues.OBJECT_KEY
                and (char == ":" or char.isspace())
            ):
                self.log(
                    "While parsing a string missing the left delimiter in object key context, we found a :, stopping here",
                )
                break
            if self.context.current == ContextValues.OBJECT_VALUE and char in [
                ",",
                "}",
            ]:
                rstring_delimiter_missing = True
                i = self.skip_to_character(character=rstring_delimiter, idx=1)
                next_c = self.get_char_at(i)
                if next_c:
                    i += 1
                    i = self.skip_whitespaces_at(idx=i, move_main_index=False)
                    next_c = self.get_char_at(i)
                    if not next_c or next_c in [",", "}"]:
                        rstring_delimiter_missing = False
                    else:

                        i = self.skip_to_character(character=lstring_delimiter, idx=i)
                        if doubled_quotes:
                            i = self.skip_to_character(
                                character=lstring_delimiter, idx=i
                            )
                        next_c = self.get_char_at(i)
                        if not next_c:
                            rstring_delimiter_missing = False
                        else:
                            i = self.skip_whitespaces_at(
                                idx=i + 1, move_main_index=False
                            )
                            next_c = self.get_char_at(i)
                            if next_c and next_c != ":":
                                rstring_delimiter_missing = False
                else:

                    i = self.skip_to_character(character=":", idx=1)
                    next_c = self.get_char_at(i)
                    if next_c:
                        break
                    else:
                        i = self.skip_whitespaces_at(idx=1, move_main_index=False)

                        j = self.skip_to_character(character="}", idx=i)
                        if j - i > 1:
                            rstring_delimiter_missing = False
                        elif self.get_char_at(j):
                            for c in reversed(string_acc):
                                if c == "{":
                                    rstring_delimiter_missing = False
                                    break
                                elif c == "}":
                                    break
                if rstring_delimiter_missing:
                    self.log(
                        "While parsing a string missing the left delimiter in object value context, we found a , or } and we couldn't determine that a right delimiter was present. Stopping here",
                    )
                    break
            if char == "]" and ContextValues.ARRAY in self.context.context:
                i = self.skip_to_character(rstring_delimiter)
                if not self.get_char_at(i):
                    break
            string_acc += char
            self.index += 1
            char = self.get_char_at()
            if char and len(string_acc) > 0 and string_acc[-1] == "\\":
                self.log("Found a stray escape sequence, normalizing it")
                if char in [rstring_delimiter, "t", "n", "r", "b", "\\"]:
                    string_acc = string_acc[:-1]
                    escape_seqs = {"t": "\t", "n": "\n", "r": "\r", "b": "\b"}
                    string_acc += escape_seqs.get(char, char) or char
                    self.index += 1
                    char = self.get_char_at()
            if (
                char == ":"
                and not missing_quotes
                and self.context.current == ContextValues.OBJECT_KEY
            ):
                i = self.skip_to_character(character=lstring_delimiter, idx=1)
                next_c = self.get_char_at(i)
                if next_c:
                    i += 1
                    i = self.skip_to_character(character=rstring_delimiter, idx=i)
                    next_c = self.get_char_at(i)
                    if next_c:
                        i += 1
                        i = self.skip_whitespaces_at(idx=i, move_main_index=False)
                        next_c = self.get_char_at(i)
                        if next_c and next_c in [",", "}"]:
                            self.log(
                                "While parsing a string missing the right delimiter in object key context, we found a :, stopping here",
                            )
                            break
                else:
                    self.log(
                        "While parsing a string missing the right delimiter in object key context, we found a :, stopping here",
                    )
                    break
            if char == rstring_delimiter:
                if doubled_quotes and self.get_char_at(1) == rstring_delimiter:
                    self.log(
                        "While parsing a string, we found a doubled quote, ignoring it"
                    )
                    self.index += 1
                elif (
                    missing_quotes
                    and self.context.current == ContextValues.OBJECT_VALUE
                ):
                    i = 1
                    next_c = self.get_char_at(i)
                    while next_c and next_c not in [
                        rstring_delimiter,
                        lstring_delimiter,
                    ]:
                        i += 1
                        next_c = self.get_char_at(i)
                    if next_c:
                        i += 1
                        i = self.skip_whitespaces_at(idx=i, move_main_index=False)
                        next_c = self.get_char_at(i)
                        if next_c and next_c == ":":
                            self.index -= 1
                            char = self.get_char_at()
                            self.log(
                                "In a string with missing quotes and object value context, I found a delimeter but it turns out it was the beginning on the next key. Stopping here.",
                            )
                            break
                elif unmatched_delimiter:
                    unmatched_delimiter = False
                    string_acc += str(char)
                    self.index += 1
                    char = self.get_char_at()
                else:
                    i = 1
                    next_c = self.get_char_at(i)
                    check_comma_in_object_value = True
                    while next_c and next_c not in [
                        rstring_delimiter,
                        lstring_delimiter,
                    ]:
                        if check_comma_in_object_value and next_c.isalpha():
                            check_comma_in_object_value = False
                        if (
                            (
                                ContextValues.OBJECT_KEY in self.context.context
                                and next_c in [":", "}"]
                            )
                            or (
                                ContextValues.OBJECT_VALUE in self.context.context
                                and next_c == "}"
                            )
                            or (
                                ContextValues.ARRAY in self.context.context
                                and next_c in ["]", ","]
                            )
                            or (
                                check_comma_in_object_value
                                and self.context.current == ContextValues.OBJECT_VALUE
                                and next_c == ","
                            )
                        ):
                            break
                        i += 1
                        next_c = self.get_char_at(i)
                    if (
                        next_c == ","
                        and self.context.current == ContextValues.OBJECT_VALUE
                    ):
                        i += 1
                        i = self.skip_to_character(character=rstring_delimiter, idx=i)
                        next_c = self.get_char_at(i)
                        i += 1
                        i = self.skip_whitespaces_at(idx=i, move_main_index=False)
                        next_c = self.get_char_at(i)
                        if next_c == "}":
                            self.log(
                                "While parsing a string, we misplaced a quote that would have closed the string but has a different meaning here since this is the last element of the object, ignoring it",
                            )
                            unmatched_delimiter = not unmatched_delimiter
                            string_acc += str(char)
                            self.index += 1
                            char = self.get_char_at()
                    elif (
                        next_c == rstring_delimiter and self.get_char_at(i - 1) != "\\"
                    ):
                        if self.context.current == ContextValues.OBJECT_VALUE:
                            i = self.skip_to_character(
                                character=rstring_delimiter, idx=i + 1
                            )
                            i += 1
                            next_c = self.get_char_at(i)
                            while next_c and next_c != ":":
                                if next_c == "," or (
                                    next_c == rstring_delimiter
                                    and self.get_char_at(i - 1) != "\\"
                                ):
                                    break
                                i += 1
                                next_c = self.get_char_at(i)
                            if next_c != ":":
                                self.log(
                                    "While parsing a string, we a misplaced quote that would have closed the string but has a different meaning here, ignoring it",
                                )
                                unmatched_delimiter = not unmatched_delimiter
                                string_acc += str(char)
                                self.index += 1
                                char = self.get_char_at()
                        elif self.context.current == ContextValues.ARRAY:
                            i = self.skip_to_character(rstring_delimiter, idx=i + 1)
                            next_c = self.get_char_at(i)
                            if next_c and next_c == rstring_delimiter:
                                i += 1
                                i = self.skip_whitespaces_at(
                                    idx=i, move_main_index=False
                                )
                                next_c = self.get_char_at(i)
                                if next_c and next_c in [",", "]"]:
                                    self.log(
                                        "While parsing a string, we a misplaced quote that would have closed the string but has a different meaning here, ignoring it",
                                    )
                                    unmatched_delimiter = not unmatched_delimiter
                                    string_acc += str(char)
                                    self.index += 1
                                    char = self.get_char_at()

        if (
            char
            and missing_quotes
            and self.context.current == ContextValues.OBJECT_KEY
            and char.isspace()
        ):
            self.log(
                "While parsing a string, handling an extreme corner case in which the LLM added a comment instead of valid string, invalidate the string and return an empty value",
            )
            self.skip_whitespaces_at()
            if self.get_char_at() not in [":", ","]:
                return ""


        if char != rstring_delimiter:
            self.log(
                "While parsing a string, we missed the closing quote, ignoring",
            )
        else:
            self.index += 1

        return string_acc.rstrip()

    def parse_number(self) -> Union[float, int, str, JSONReturnType]:
        number_str = ""
        number_chars = set("0123456789-.eE/,")
        char = self.get_char_at()
        is_array = self.context.current == ContextValues.ARRAY
        while char and char in number_chars and (char != "," or not is_array):
            number_str += char
            self.index += 1
            char = self.get_char_at()
        if len(number_str) > 1 and number_str[-1] in "-eE/,":
            number_str = number_str[:-1]
            self.index -= 1
        try:
            if "," in number_str:
                return str(number_str)
            if "." in number_str or "e" in number_str or "E" in number_str:
                return float(number_str)
            elif number_str == "-":
                return self.parse_json()
            else:
                return int(number_str)
        except ValueError:
            return number_str

    def parse_boolean_or_null(self) -> Union[bool, str, None]:
        starting_index = self.index
        char = (self.get_char_at() or "").lower()
        value: Optional[Tuple[str, Optional[bool]]]
        if char == "t":
            value = ("true", True)
        elif char == "f":
            value = ("false", False)
        elif char == "n":
            value = ("null", None)

        if value:
            i = 0
            while char and i < len(value[0]) and char == value[0][i]:
                i += 1
                self.index += 1
                char = (self.get_char_at() or "").lower()
            if i == len(value[0]):
                return value[1]

        self.index = starting_index
        return ""

    def get_char_at(self, count: int = 0) -> Union[str, Literal[False]]:
        try:
            return self.json_str[self.index + count]
        except IndexError:
            return False

    def skip_whitespaces_at(self, idx: int = 0, move_main_index=True) -> int:

        try:
            char = self.json_str[self.index + idx]
        except IndexError:
            return idx
        while char.isspace():
            if move_main_index:
                self.index += 1
            else:
                idx += 1
            try:
                char = self.json_str[self.index + idx]
            except IndexError:
                return idx
        return idx

    def skip_to_character(self, character: str, idx: int = 0) -> int:

        try:
            char = self.json_str[self.index + idx]
        except IndexError:
            return idx
        while char != character:
            idx += 1
            try:
                char = self.json_str[self.index + idx]
            except IndexError:
                return idx
        if self.index + idx > 0 and self.json_str[self.index + idx - 1] == "\\":
            return self.skip_to_character(character=character, idx=idx + 1)
        return idx

    def _log(self, text: str) -> None:
        window: int = 10
        start: int = max(self.index - window, 0)
        end: int = min(self.index + window, len(self.json_str))
        context: str = self.json_str[start:end]
        self.logger.append(
            {
                "text": text,
                "context": context,
            }
        )



import argparse
import json
import sys
from typing import Dict, List, Optional, TextIO, Tuple, Union



def repair_json(
    json_str: str = "",
    return_objects: bool = False,
    skip_json_loads: bool = False,
    logging: bool = False,
    json_fd: Optional[TextIO] = None,
    ensure_ascii: bool = True,
    chunk_length: int = 0,
) -> Union[JSONReturnType, Tuple[JSONReturnType, List[Dict[str, str]]]]:

    parser = JSONParser(json_str, json_fd, logging, chunk_length)
    if skip_json_loads:
        parsed_json = parser.parse()
    else:
        try:
            if json_fd:
                parsed_json = json.load(json_fd)
            else:
                parsed_json = json.loads(json_str)
        except json.JSONDecodeError:
            parsed_json = parser.parse()
    if return_objects or logging:
        return parsed_json
    return json.dumps(parsed_json, ensure_ascii=ensure_ascii)


def loads(
    json_str: str,
    skip_json_loads: bool = False,
    logging: bool = False,
) -> Union[JSONReturnType, Tuple[JSONReturnType, List[Dict[str, str]]]]:

    return repair_json(
        json_str=json_str,
        return_objects=True,
        skip_json_loads=skip_json_loads,
        logging=logging,
    )


def load(
    fd: TextIO,
    skip_json_loads: bool = False,
    logging: bool = False,
    chunk_length: int = 0,
) -> Union[JSONReturnType, Tuple[JSONReturnType, List[Dict[str, str]]]]:

    return repair_json(
        json_fd=fd,
        chunk_length=chunk_length,
        return_objects=True,
        skip_json_loads=skip_json_loads,
        logging=logging,
    )


def from_file(
    filename: str,
    skip_json_loads: bool = False,
    logging: bool = False,
    chunk_length: int = 0,
) -> Union[JSONReturnType, Tuple[JSONReturnType, List[Dict[str, str]]]]:

    with open(filename) as fd:
        jsonobj = load(
            fd=fd,
            skip_json_loads=skip_json_loads,
            logging=logging,
            chunk_length=chunk_length,
        )

    return jsonobj



if __name__ == "__main__":
    bad_json = """
    {
        "key": 'value',
        "key2": value2,
        "key3": "value3",
        }
    """
    json = repair_json(bad_json )
    print(json , type(json))